/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#ifndef CATLASS_GEMM_KERNEL_MATMUL_FAULT_TOLERANCE_HPP
#define CATLASS_GEMM_KERNEL_MATMUL_FAULT_TOLERANCE_HPP

#include "catlass/catlass.hpp"
#include "catlass/arch/resource.hpp"
#include "catlass/arch/cross_core_sync.hpp"
#include "catlass/layout/layout.hpp"
#include "catlass/matrix_coord.hpp"
#include "catlass/gemm_coord.hpp"
#include "catlass/gemv_coord.hpp"
#include "catlass/gemv/helper.hpp"

namespace Catlass::Gemm::Kernel
{

    template <
        class BlockMmad_,
        class BlockSumGemvTransposeA_,
        class BlockSumGemvB_,
        class BlockSumGemvTransposeC_,
        class BlockSumGemvC_,
        class BlockScheduler_>
    class MatmulFaultTolerance
    {
    public:
        using BlockMmad = BlockMmad_;
        using BlockSumGemvTransposeA = BlockSumGemvTransposeA_;
        using BlockSumGemvB = BlockSumGemvB_;
        using BlockSumGemvTransposeC = BlockSumGemvTransposeC_;
        using BlockSumGemvC = BlockSumGemvC_;
        using ArchTag = typename BlockMmad::ArchTag;
        using L1TileShape = typename BlockMmad::L1TileShape;
        using ElementA = typename BlockMmad::ElementA;
        using LayoutA = typename BlockMmad::LayoutA;
        using ElementB = typename BlockMmad::ElementB;
        using LayoutB = typename BlockMmad::LayoutB;
        using ElementC = typename BlockMmad::ElementC;
        using LayoutC = typename BlockMmad::LayoutC;

        using LayoutTransposeA = typename BlockSumGemvTransposeA::LayoutA;
        using UBTileShapeTransposeA = typename BlockSumGemvTransposeA::UBTileShape;

        using ElementX = typename BlockSumGemvB::ElementX;
        using LayoutX = typename BlockSumGemvB::LayoutX;
        using ElementY = typename BlockSumGemvB::ElementY;
        using LayoutY = typename BlockSumGemvB::LayoutY;
        using UBTileShapeB = typename BlockSumGemvB::UBTileShape;

        using LayoutTransposeC = typename BlockSumGemvTransposeC::LayoutA;
        using UBTileShapeTransposeC = typename BlockSumGemvTransposeC::UBTileShape;

        using UBTileShapeC = typename BlockSumGemvC::UBTileShape;

        using UBAlignHelper = Gemv::helper::UBAlignHelper<ElementA>;

        using ElementD = ElementX;

        using BlockScheduler = BlockScheduler_;

        static_assert(std::is_same_v<typename BlockSumGemvTransposeA::ElementA, ElementA> &&
                          std::is_same_v<typename BlockSumGemvB::ElementA, ElementA> &&
                          std::is_same_v<ElementA, ElementB>,
                      "The AType of Gemv and AType and BType of Mmad should be consistent.");

        static_assert(std::is_same_v<ElementX, ElementY> && std::is_same_v<ElementX, ElementD>,
                      "The XType and YType of Gemv and DType should be consistent.");

        static_assert(std::is_same_v<LayoutA, LayoutB>,
                      "The LayoutA and LayoutB of Gemv should be consistent.");

        enum class AivCore
        {
            AIV0 = 0,
            AIV1
        };

        enum class CheckSumType
        {
            NO_CHECKSUM = 0,
            COL_CHECKSUM,
            ROW_CHECKSUM,
            ALL_CHECKSUM
        };

        /// Parameters structure
        struct Params
        {
            // Data members
            GemmCoord problemShape;
            GM_ADDR ptrA;
            LayoutA layoutA;
            GM_ADDR ptrB;
            LayoutB layoutB;
            GM_ADDR ptrC;
            LayoutC layoutC;
            GM_ADDR ptrD;
            CheckSumType checkSumType;
            GM_ADDR ptrWorkspace;

            // Methods
            CATLASS_HOST_DEVICE
            Params() {}

            CATLASS_HOST_DEVICE
            Params(
                GemmCoord const &problemShape_,
                GM_ADDR ptrA_, LayoutA const &layoutA_,
                GM_ADDR ptrB_, LayoutB const &layoutB_,
                GM_ADDR ptrC_, LayoutC const &layoutC_,
                GM_ADDR ptrD_, int32_t checkSumType_, GM_ADDR ptrWorkspace_) : problemShape(problemShape_), ptrA(ptrA_), layoutA(layoutA_), ptrB(ptrB_), layoutB(layoutB_),
                                                                               ptrC(ptrC_), layoutC(layoutC_), ptrD(ptrD_), checkSumType(static_cast<CheckSumType>(checkSumType_)), ptrWorkspace(ptrWorkspace_) {}
        };

        struct Arguments
        {
            GemmCoord problemShape;
            GM_ADDR ptrA;
            GM_ADDR ptrB;
            GM_ADDR ptrC;
            GM_ADDR ptrD;
            int32_t checkSumType;
        };

        static bool CanImplement(const Arguments &args)
        {
            return true;
        }

        static size_t GetWorkspaceSize(const Arguments &args)
        {
            return 0;
        }

        static Params ToUnderlyingArguments(const Arguments &args, uint8_t *workspace)
        {
            GemmCoord problemShape = args.problemShape;
            uint32_t m = problemShape.m();
            uint32_t n = problemShape.n();
            uint32_t k = problemShape.k();
            LayoutA layoutA{m, k};
            LayoutB layoutB{k, n};
            LayoutC layoutC{m, n};
            Params params{problemShape, args.ptrA, layoutA, args.ptrB, layoutB, args.ptrC, layoutC, args.ptrD, args.checkSumType, workspace};
            return params;
        }

        // Methods
        CATLASS_DEVICE
        MatmulFaultTolerance() {}

        template <int32_t CORE_TYPE = g_coreType>
        CATLASS_DEVICE void operator()(Params const &params);

        template <>
        CATLASS_DEVICE void operator()<AscendC::AIC>(Params const &params)
        {
            BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1TileShape::M, L1TileShape::N));
            uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();

            BlockMmad blockMmad(resource);

            // Represent the full gm
            AscendC::GlobalTensor<ElementA> gmA;
            gmA.SetGlobalBuffer((__gm__ ElementA *)params.ptrA);
            AscendC::GlobalTensor<ElementB> gmB;
            gmB.SetGlobalBuffer((__gm__ ElementB *)params.ptrB);
            AscendC::GlobalTensor<ElementC> gmC;
            gmC.SetGlobalBuffer((__gm__ ElementC *)params.ptrC);

            for (uint32_t loopIdx = AscendC::GetBlockIdx(); loopIdx < coreLoops; loopIdx += AscendC::GetBlockNum())
            {
                // Compute block location
                GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);

                // Compute initial location in logical coordinates
                MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, blockCoord.k() * L1TileShape::K};
                MatrixCoord offsetB{blockCoord.k() * L1TileShape::K, blockCoord.n() * L1TileShape::N};
                MatrixCoord offsetC{blockCoord.m() * L1TileShape::M, blockCoord.n() * L1TileShape::N};
                int64_t gmOffsetA = params.layoutA.GetOffset(offsetA);
                int64_t gmOffsetB = params.layoutB.GetOffset(offsetB);
                int64_t gmOffsetC = params.layoutC.GetOffset(offsetC);

                // Compute block-scoped matrix multiply-add
                blockMmad(
                    gmA[gmOffsetA], params.layoutA,
                    gmB[gmOffsetB], params.layoutB,
                    gmC[gmOffsetC], params.layoutC,
                    actualBlockShape);

                // Synchronize cross core
                Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(flagAicFinishStore);
            }

            AscendC::PipeBarrier<PIPE_ALL>();
        }

        template <>
        CATLASS_DEVICE void operator()<AscendC::AIV>(Params const &params)
        {
            AscendC::SetAtomicNone();

            BlockScheduler matmulBlockScheduler(params.problemShape, MakeCoord(L1TileShape::M, L1TileShape::N));
            uint32_t coreLoops = matmulBlockScheduler.GetCoreLoops();

            // Represent the full gm
            AscendC::GlobalTensor<ElementA> gmA;
            gmA.SetGlobalBuffer((__gm__ ElementA *)params.ptrA);
            AscendC::GlobalTensor<ElementB> gmB;
            gmB.SetGlobalBuffer((__gm__ ElementB *)params.ptrB);
            AscendC::GlobalTensor<ElementC> gmC;
            gmC.SetGlobalBuffer((__gm__ ElementC *)params.ptrC);
            AscendC::GlobalTensor<ElementD> gmD;
            gmD.SetGlobalBuffer((__gm__ ElementD *)params.ptrD);

            // Get aicore information
            uint32_t aicoreIndex = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
            // AscendC::printf("the AscendC::GetBlockIdx() is : %d \n", AscendC::GetBlockIdx());
            // AscendC::printf("the AscendC::GetSubBlockNum() is : %d \n", AscendC::GetSubBlockNum());

            uint32_t aicoreNum = AscendC::GetBlockNum();
            // AscendC::printf("the AscendC::GetBlockNum() is : %d \n", AscendC::GetBlockNum());
            AivCore aivCore = static_cast<AivCore>(AscendC::GetSubBlockIdx());

            MatrixCoord loopsMN = matmulBlockScheduler.loopsMN;
            uint32_t maxMN = max(loopsMN.row(), loopsMN.column());

            if (params.checkSumType == CheckSumType::ALL_CHECKSUM)
            { // NOT SUPPORTED YET
                if (aivCore == AivCore::AIV0)
                {
                    uint32_t UBTileMRound = RoundUp(UBTileShapeTransposeA::M, UBAlignHelper::ALIGN);
                    uint32_t UBTileNRound = RoundUp(UBTileShapeTransposeA::N, UBAlignHelper::ALIGN);
                    BlockSumGemvTransposeA blockSumGemv(resource);
                    for (uint32_t loopIdx = aicoreIndex; loopIdx < loopsMN.row(); loopIdx += aicoreNum)
                    {
                        uint32_t mActual = (loopIdx == (loopsMN.row() - 1)) ? (params.problemShape.m() - loopIdx * L1TileShape::M) : L1TileShape::M;

                        uint32_t Nsum = CeilDiv(params.problemShape.k(), UBTileMRound);
                        for (uint32_t sumId = 0; sumId < Nsum; ++sumId)
                        {
                            MatrixCoord offsetA{loopIdx * L1TileShape::M, sumId * UBTileMRound};
                            int64_t gmOffsetA = params.layoutA.GetOffset(offsetA);
                            int64_t gmOffsetEA = loopIdx * params.problemShape.k() + sumId * UBTileMRound;
                            uint32_t kActual = (sumId == Nsum - 1) ? (params.problemShape.k() - sumId * UBTileMRound) : UBTileMRound;
                            GemvCoord actualBlockShape = GemvCoord{kActual, mActual};
                            layout::VectorLayout layoutEA{kActual};
                            LayoutTransposeA layoutTransposeA(params.layoutA.shape(1), params.layoutA.shape(0));
                            blockSumGemv.rowSum(gmA[gmOffsetA], layoutTransposeA, gmD[gmOffsetEA], layoutEA, actualBlockShape);
                        }
                    }

                    AscendC::SyncAll<true>();

                    // Loop through the epilogue calculations of each basic block
                    for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum)
                    {
                        // Compute block location
                        GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                        GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);

                        // Compute initial location in logical coordinates
                        MatrixCoord offsetB{0, blockCoord.n() * L1TileShape::N};
                        int64_t gmOffsetB = params.layoutB.GetOffset(offsetB);
                        uint32_t gmoffsetEA = blockCoord.m() * params.problemShape.k();
                        uint32_t gmoffsetEAB = (loopsMN.row() + loopsMN.column()) * params.problemShape.k() + loopIdx * L1TileShape::N;
                        layout::VectorLayout layoutEA{params.problemShape.k()};
                        layout::VectorLayout layoutEAB{L1TileShape::N};
                        GemvCoord gemvActualBlockShape = GemvCoord{actualBlockShape.n(), params.problemShape.k()};
                        LayoutTransposeA layoutTransposeB(params.layoutB.shape(1), params.layoutB.shape(0));
                        blockSumGemv(gmB[gmOffsetB], layoutTransposeB, gmD[gmoffsetEA], layoutEA, gmD[gmoffsetEAB], layoutEAB, gemvActualBlockShape);

                        // Synchronize cross core
                        Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);
                    }
                }
                else if (aivCore == AivCore::AIV1)
                {
                    uint32_t UBTileMRound = RoundUp(UBTileShapeB::M, UBAlignHelper::ALIGN);
                    uint32_t UBTileNRound = RoundUp(UBTileShapeB::N, UBAlignHelper::ALIGN);
                    BlockSumGemvB blockSumGemv(resource);

                    for (uint32_t loopIdx = aicoreIndex; loopIdx < loopsMN.column(); loopIdx += aicoreNum)
                    {
                        uint32_t nActual = (loopIdx == (loopsMN.column() - 1)) ? (params.problemShape.n() - loopIdx * L1TileShape::N) : L1TileShape::N;

                        uint32_t Nsum = CeilDiv(params.problemShape.k(), UBTileMRound);
                        for (uint32_t sumId = 0; sumId < Nsum; ++sumId)
                        {
                            MatrixCoord offsetB{sumId * UBTileMRound, loopIdx * L1TileShape::N};
                            int64_t gmOffsetB = params.layoutB.GetOffset(offsetB);
                            int64_t gmOffsetBE = (loopsMN.row() + loopIdx) * params.problemShape.k() + sumId * UBTileMRound;
                            uint32_t kActual = (sumId == Nsum - 1) ? (params.problemShape.k() - sumId * UBTileMRound) : UBTileMRound;
                            GemvCoord actualBlockShape = GemvCoord{kActual, nActual};
                            layout::VectorLayout layoutBE{kActual};
                            blockSumGemv.rowSum(gmB[gmOffsetB], params.layoutB, gmD[gmOffsetBE], layoutBE, actualBlockShape);
                        }
                    }

                    AscendC::SyncAll<true>();

                    // Loop through the epilogue calculations of each basic block
                    for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum)
                    {
                        // Compute block location
                        GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                        GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);

                        // Compute initial location in logical coordinates
                        MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, 0};
                        int64_t gmOffsetA = params.layoutA.GetOffset(offsetA);
                        uint32_t gmoffsetBE = (loopsMN.row() + blockCoord.n()) * params.problemShape.k();
                        uint32_t gmoffsetABE = (loopsMN.row() + loopsMN.column()) * params.problemShape.k() + coreLoops * L1TileShape::N + loopIdx * L1TileShape::M;
                        layout::VectorLayout layoutBE{params.problemShape.k()};
                        layout::VectorLayout layoutABE{L1TileShape::M};
                        GemvCoord gemvActualBlockShape = GemvCoord{actualBlockShape.m(), params.problemShape.k()};
                        blockSumGemv(gmA[gmOffsetA], params.layoutA, gmD[gmoffsetBE], layoutBE, gmD[gmoffsetABE], layoutABE, gemvActualBlockShape);

                        // Synchronize cross core
                        Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);
                    }
                }
            }
            else if (params.checkSumType == CheckSumType::COL_CHECKSUM)
            {
                if (aivCore == AivCore::AIV0)
                {
                    uint32_t UBTileMRound = RoundUp(UBTileShapeTransposeA::M, UBAlignHelper::ALIGN);
                    uint32_t UBTileNRound = RoundUp(UBTileShapeTransposeA::N, UBAlignHelper::ALIGN);
                    BlockSumGemvTransposeA blockSumGemv(resource);
                    for (uint32_t loopIdx = aicoreIndex; loopIdx < loopsMN.row(); loopIdx += aicoreNum)
                    {
                        uint32_t mActual = (loopIdx == (loopsMN.row() - 1)) ? (params.problemShape.m() - loopIdx * L1TileShape::M) : L1TileShape::M;

                        uint32_t Nsum = CeilDiv(params.problemShape.k(), UBTileMRound);
                        for (uint32_t sumId = 0; sumId < Nsum; ++sumId)
                        {
                            MatrixCoord offsetA{loopIdx * L1TileShape::M, sumId * UBTileMRound};
                            int64_t gmOffsetA = params.layoutA.GetOffset(offsetA);
                            int64_t gmOffsetEA = loopIdx * params.problemShape.k() + sumId * UBTileMRound;
                            uint32_t kActual = (sumId == Nsum - 1) ? (params.problemShape.k() - sumId * UBTileMRound) : UBTileMRound;
                            GemvCoord actualBlockShape = GemvCoord{kActual, mActual};
                            layout::VectorLayout layoutEA{kActual};
                            LayoutTransposeA layoutTransposeA(params.layoutA.shape(1), params.layoutA.shape(0));
                            blockSumGemv.rowSum(gmA[gmOffsetA], layoutTransposeA, gmD[gmOffsetEA], layoutEA, actualBlockShape);
                        }
                    }

                    AscendC::SyncAll<true>();

                    // Loop through the epilogue calculations of each basic block
                    for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum)
                    {
                        // Compute block location
                        GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                        GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);

                        // Compute initial location in logical coordinates
                        MatrixCoord offsetB{0, blockCoord.n() * L1TileShape::N};
                        int64_t gmOffsetB = params.layoutB.GetOffset(offsetB);
                        uint32_t gmoffsetEA = blockCoord.m() * params.problemShape.k();
                        uint32_t gmoffsetEAB = loopsMN.row() * params.problemShape.k() + loopIdx * L1TileShape::N;
                        layout::VectorLayout layoutEA{params.problemShape.k()};
                        layout::VectorLayout layoutEAB{L1TileShape::N};
                        GemvCoord gemvActualBlockShape = GemvCoord{actualBlockShape.n(), params.problemShape.k()};
                        LayoutTransposeA layoutTransposeB(params.layoutB.shape(1), params.layoutB.shape(0));
                        blockSumGemv(gmB[gmOffsetB], layoutTransposeB, gmD[gmoffsetEA], layoutEA, gmD[gmoffsetEAB], layoutEAB, gemvActualBlockShape);

                        // Synchronize cross core
                        Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);
                    }
                }
                else if (aivCore == AivCore::AIV1)
                {
                    uint32_t UBTileMRound = RoundUp(UBTileShapeC::M, UBAlignHelper::ALIGN);
                    uint32_t UBTileNRound = RoundUp(UBTileShapeC::N, UBAlignHelper::ALIGN);
                    BlockSumGemvTransposeC blockSumGemv(resource);

                    AscendC::SyncAll<true>();

                    for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum)
                    {
                        // Compute block location
                        GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                        GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
                        // Synchronize cross core
                        Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);

                        MatrixCoord offsetC{blockCoord.m() * L1TileShape::M, blockCoord.n() * L1TileShape::N};
                        int64_t gmOffsetC = params.layoutC.GetOffset(offsetC);
                        int64_t gmOffsetEC = loopsMN.row() * params.problemShape.k() + coreLoops * L1TileShape::N + loopIdx * L1TileShape::N;
                        layout::VectorLayout layoutEC{actualBlockShape.n()};
                        GemvCoord gemvActualBlockShape = GemvCoord{actualBlockShape.n(), actualBlockShape.m()};
                        LayoutTransposeC layoutTransposeC(params.layoutC.shape(1), params.layoutC.shape(0));
                        blockSumGemv.rowSum(gmC[gmOffsetC], layoutTransposeC, gmD[gmOffsetEC], layoutEC, gemvActualBlockShape);
                    }
                }
            }
            else if (params.checkSumType == CheckSumType::ROW_CHECKSUM)
            {
                if (aivCore == AivCore::AIV0)
                {
                    uint32_t UBTileMRound = RoundUp(UBTileShapeC::M, UBAlignHelper::ALIGN);
                    uint32_t UBTileNRound = RoundUp(UBTileShapeC::N, UBAlignHelper::ALIGN);
                    BlockSumGemvC blockSumGemv(resource);

                    AscendC::SyncAll<true>();

                    for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum)
                    {
                        // Compute block location
                        GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                        GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);
                        // Synchronize cross core
                        Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);

                        MatrixCoord offsetC{blockCoord.m() * L1TileShape::M, blockCoord.n() * L1TileShape::N};
                        int64_t gmOffsetC = params.layoutC.GetOffset(offsetC);
                        int64_t gmOffsetCE = loopsMN.column() * params.problemShape.k() + coreLoops * L1TileShape::M + loopIdx * L1TileShape::M;
                        layout::VectorLayout layoutCE{actualBlockShape.m()};
                        GemvCoord gemvActualBlockShape = GemvCoord{actualBlockShape.m(), actualBlockShape.n()};
                        AscendC::DumpTensor(gmC[gmOffsetC], coreLoops, L1TileShape::M);
                        blockSumGemv.rowSum(gmC[gmOffsetC], params.layoutC, gmD[gmOffsetCE], layoutCE, gemvActualBlockShape);
                        AscendC::DumpTensor(gmD[gmOffsetCE], coreLoops + 1, L1TileShape::M);
                    }
                }
                else if (aivCore == AivCore::AIV1)
                {
                    uint32_t UBTileMRound = RoundUp(UBTileShapeB::M, UBAlignHelper::ALIGN);
                    uint32_t UBTileNRound = RoundUp(UBTileShapeB::N, UBAlignHelper::ALIGN);
                    BlockSumGemvB blockSumGemv(resource);

                    for (uint32_t loopIdx = aicoreIndex; loopIdx < loopsMN.column(); loopIdx += aicoreNum)
                    {
                        uint32_t nActual = (loopIdx == (loopsMN.column() - 1)) ? (params.problemShape.n() - loopIdx * L1TileShape::N) : L1TileShape::N;

                        uint32_t Nsum = CeilDiv(params.problemShape.k(), UBTileMRound);
                        for (uint32_t sumId = 0; sumId < Nsum; ++sumId)
                        {
                            MatrixCoord offsetB{sumId * UBTileMRound, loopIdx * L1TileShape::N};
                            int64_t gmOffsetB = params.layoutB.GetOffset(offsetB);
                            int64_t gmOffsetBE = loopIdx * params.problemShape.k() + sumId * UBTileMRound;
                            uint32_t kActual = (sumId == Nsum - 1) ? (params.problemShape.k() - sumId * UBTileMRound) : UBTileMRound;
                            GemvCoord actualBlockShape = GemvCoord{kActual, nActual};
                            layout::VectorLayout layoutBE{kActual};
                            blockSumGemv.rowSum(gmB[gmOffsetB], params.layoutB, gmD[gmOffsetBE], layoutBE, actualBlockShape);
                        }
                    }

                    AscendC::SyncAll<true>();

                    // Loop through the epilogue calculations of each basic block
                    for (uint32_t loopIdx = aicoreIndex; loopIdx < coreLoops; loopIdx += aicoreNum)
                    {
                        // Compute block location
                        GemmCoord blockCoord = matmulBlockScheduler.GetBlockCoord(loopIdx);
                        GemmCoord actualBlockShape = matmulBlockScheduler.GetActualBlockShape(blockCoord);

                        // Compute initial location in logical coordinates
                        MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, 0};
                        int64_t gmOffsetA = params.layoutA.GetOffset(offsetA);
                        uint32_t gmoffsetBE = blockCoord.n() * params.problemShape.k();
                        uint32_t gmoffsetABE = loopsMN.column() * params.problemShape.k() + loopIdx * L1TileShape::M;
                        layout::VectorLayout layoutBE{params.problemShape.k()};
                        layout::VectorLayout layoutABE{L1TileShape::M};
                        GemvCoord gemvActualBlockShape = GemvCoord{actualBlockShape.m(), params.problemShape.k()};
                        blockSumGemv(gmA[gmOffsetA], params.layoutA, gmD[gmoffsetBE], layoutBE, gmD[gmoffsetABE], layoutABE, gemvActualBlockShape);

                        // Synchronize cross core
                        Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagAicFinishStore);
                    }
                }
            }
        }

    private:
        // ID used for inter-core synchronization
        static constexpr Arch::FlagID FLAG_AIC_FINISH_STORE = 0;
        static constexpr Arch::FlagID RV_FLAG_AIC_FINISH_STORE = 1;
        Arch::CrossCoreFlagWithReverse<> flagAicFinishStore{FLAG_AIC_FINISH_STORE, RV_FLAG_AIC_FINISH_STORE};
        Arch::Resource<ArchTag> resource;
    };

} // namespace Catlass::Gemm::Kernel

#endif // CATLASS_GEMM_KERNEL_MATMUL_FAULT_TOLERANCE_HPP
